# generatedvh-test.py

# Test code to generate a dDVH from a DICOM-RT file(s).
# Roy Keyes (roy.coding@gmail)
# 4 Mar 2010

import dicom
import numpy as np
import pylab as pl
import matplotlib.nxutils as nx
import dvhdoses # Just for debugging


def gendvh(contour, thickness, ss, rtd, plan):
    '''Example of dDVH generatation from RTSS and RTDOSE files.'''

    # Get list of contour slices
    cslist = ss.ROIContours[contour].Contours

    # Convert contours to sets of triples
    csa = []
    for x in cslist:
        j = 0
        jmax = len(x.ContourData)
        ca=[]
        while j < jmax:
            ca += [x.ContourData[j:j+3]]
            j += 3
        
        csa += [np.array(ca)]

    # Dose data from pixel array
    dose = rtd.pixel_array

    # Make coordinate array
    xspace,yspace = rtd.PixelSpacing # mm spacing bewteen points
    zspace = rtd.SliceThickness
    # SliceThickness isn't always there...
    if not zspace:
        print 'RTDose SliceThickness missing'
        zspace = thickness # mm slice width <-- Should be in Dicom file...
    else:
        zspace = float(zspace)
    ashape = dose.shape

    # Image position offset
    pos = rtd.ImagePositionPatient # mm

    ###NOTE####
    # Coordinate order is not uniform between dose, contour, offset, etc arrays
    # Therefore he we use dose x for contour y and vice versa...
    ###########

    xoffset = pos[0]
    yoffset = pos[1]
    zoffset = pos[2]

    coord = []
    j,k = 0,0
    while j < ashape[2]:
        k = 0
        while k < ashape[1]:
            coord += [[j*xspace + xoffset,k*yspace + yoffset]]
            k += 1
        j += 1

    coord = np.array(coord)

    # Filter function to use with map to keep only dose values inside the contour
    def filt(a,b):
        if a:
            return b
        else:
            return []

    # Make lists of points inside contours and filter them
    # Simultaneously determine max dose inside contour
    rplist = []
    cmax = 0.0
    #cmin = []
    for x in csa:
        # Exclude z (transverse) dimension from contour array
        inside = nx.points_inside_poly(coord, x[:,0:2])

        # Determine which slice contour points belong to
        zslice = int(abs(zoffset - x[0][2])/zspace)

        # Unravel dose array to match coord array
        # Dose must first be transposed to reorder indices to match
        #   coordinate array index order
        roipoints = filter(lambda y:y,map(filt,inside,dose[zslice].T.reshape(ashape[1]*ashape[2])))
        #print zslice,'len(roipoints):',len(roipoints), 'ashape[1]*ashape[2]:',ashape[1]*ashape[2]

        # Get max dose in contour for bin limits
        #print 'DEBUG:', np.array(roipoints)*(100)*1.4e-05
        try:
            tmpmax = max(roipoints)
            #tmpmin = min(roipoints)
        except ValueError:
            pass
        if cmax < tmpmax:
            cmax = tmpmax
        #cmin += [tmpmin]

        # Determine voxel volume for given slice
        if zslice == 0 or zslice == (dose.shape[0] -1):
            vvol = 0.5*xspace*yspace*zspace
        else:
            vvol = xspace*yspace*zspace

        rplist += [(vvol,roipoints)]

    #cmin = min(cmin)
    #print 'cmin:',cmin

    # Bin the dose in each contour/slice 
    # Prescription dose (Gy)
    rx = plan.DoseReferences[0].TargetPrescriptionDose
    print 'Prescription Dose:',rx*100,'cGy'

    # Dose scaling factor
    dosescale = rtd.DoseGridScaling
    print 'Dose Grid Scaling (to Gy):', dosescale

    # Max and min doses (scaled to correct units)   
    dmax = cmax*dosescale
    #dmin = cmin*dosescale
    print 'Max dose:', dmax*100,'(cGy)'
    #print 'Min dose:', dmin*100,'(cGy)'

    # Array of bins to use
    bins = int(dmax*100 + 1)
    print 'Bins (cGy):',bins

    histo = np.zeros(bins)
    totvol = 0.0
    for x,y in rplist:      
        if len(y) == 0:
            continue
        else:
            # Convert doses to cGy
            y = np.array(y)*(100)*1.4e-05
            counts,b = np.histogram(y,bins=bins,range=(0,bins))

            histo += x*counts  
            totvol += x*len(y)
    
    #print histo
    print 'Total Volume:',totvol,'(mm^3)'

    return histo,totvol


def calccdvh(ddvh):
    '''Calculate cDVH from dDVH.'''
    # Scratch code for now...
    # cDVH(x) is Sum (Integral) of dDVH with x as lower limit
    cdvh = []
    j = 0
    jmax = len(ddvh)
    while j < jmax:
        cdvh += [sum(ddvh[j:])]
        j += 1

    cdvh = np.array(cdvh)
    return cdvh

todo = '''
1. Verify DVH correctness
2. Test other contours
3. Clean up code for inclusion
'''

print todo


ss = dicom.read_file('rtss.dcm')
rtd = dicom.read_file('rtdose.dcm')
plan  = dicom.read_file('rtplan.dcm')

# Set contour number
contour = 4
thickness = 3.0 #mm, this was determined from the contours, 
                # but should be in the dose file
histo,vol = gendvh(contour, thickness, ss, rtd, plan)

#print histo

cdvh = np.array(calccdvh(histo))

#print cdvh

####
# Get DVH from dicom file for comparison
dicomdvh = np.array(rtd.DVHs[3].DVHData)[1::2]
# Dicom dDVH
dicomddvh = np.array(dvhdoses.get_ddvh(dicomdvh))


## Add light grey grid
ax=pl.axes()
ax.set_axisbelow(True)
pl.grid(color='lightgrey', zorder=20)

#pl.plot(histo)
#pl.plot(100*histo/histo.max(), label="calculated")
#pl.plot(100*dicomddvh/dicomddvh.max(),label="dicom")

pl.plot(100*cdvh/cdvh.max(),label="calculated")
pl.plot(100*dicomdvh/dicomdvh.max(),label="dicom")

pl.legend()

pl.show()


